import itertools
import pandas as pd
from sklearn.datasets import load_iris
from sklearn.preprocessing import KBinsDiscretizer

# 步骤 2：生成候选项集
def create_candidates(data, k):
    candidates = set()
    for transaction in data:
        for item in itertools.combinations(sorted(transaction), k):
            candidates.add(item)
    return list(candidates)

# 步骤 3：计算支持度
def compute_support(data, candidates):
    support = {}
    total_transactions = len(data)
    for candidate in candidates:
        count = sum(1 for transaction in data if set(candidate).issubset(set(transaction)))
        support[candidate] = count / total_transactions
    return support

# 步骤 4：生成频繁项集
def apriori(data, min_support):
    k = 1
    frequent_itemsets = []

    # 初始候选项集为单个商品的集合
    candidates = create_candidates(data, k)
    support = compute_support(data, candidates)

    # 过滤出频繁项集
    frequent_itemsets_k = {item: sup for item, sup in support.items() if sup >= min_support}
    frequent_itemsets.append(frequent_itemsets_k)

    # 不断增加项集的大小
    while frequent_itemsets_k:
        k += 1
        candidates = create_candidates([list(item) for item in frequent_itemsets_k], k)
        support = compute_support(data, candidates)
        frequent_itemsets_k = {item: sup for item, sup in support.items() if sup >= min_support}

        if frequent_itemsets_k:
            frequent_itemsets.append(frequent_itemsets_k)

    return frequent_itemsets

# 步骤 5：预处理Iris数据并调用Apriori算法
if __name__ == "__main__":
    # 加载Iris数据集
    iris = load_iris()
    df = pd.DataFrame(data=iris.data, columns=iris.feature_names)
    df['species'] = iris.target

    # 离散化数值型特征
    discretizer = KBinsDiscretizer(n_bins=3, encode='ordinal', strategy='uniform')
    df_discretized = pd.DataFrame(discretizer.fit_transform(df[df.columns[:-1]]))

    # 将离散化的特征值转换为字符串形式，方便Apriori算法处理
    for col in df_discretized.columns:
        df_discretized[col] = df_discretized[col].apply(lambda x: f'{df.columns[col]}_{int(x)}')

    # 将数据转换为事务列表的形式
    transactions = df_discretized.apply(lambda x: x.tolist(), axis=1).tolist()

    min_support = 0.2 # 设置最小支持度
    frequent_itemsets = apriori(transactions, min_support)

    # 输出结果
    for k, itemsets in enumerate(frequent_itemsets, start=1):
        print(f"频繁{k}-项集:")
        for itemset, support in itemsets.items():
            print(f"  {itemset}: 支持度 = {support:.2f}")